from enum import IntEnum

from cryptography.hazmat import backends
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.ciphers import algorithms, Cipher, modes
from typing_extensions import Buffer


class EncryptionAlgorithm(IntEnum):
    Plaintext = 0
    AES128CTR = 1
    AES128CBC = 2
    AES256CTR = 3
    AES256CBC = 4

    @property
    def key_size(self):
        """
        Returns the key size of the algorithm in bytes
        """
        if self == EncryptionAlgorithm.Plaintext:
            return 0
        elif self == EncryptionAlgorithm.AES128CBC or self == EncryptionAlgorithm.AES128CTR:
            return 16
        return 32

    @property
    def nonce_iv_size(self):
        """
        Returns the size of the nonce/IV in bytes
        """
        if self == EncryptionAlgorithm.Plaintext:
            return 0
        return 16

    @property
    def is_block_cipher(self) -> bool:
        """
        Returns whether the algorithm is a block cipher
        """
        if self == EncryptionAlgorithm.Plaintext:
            return False
        return True

    @property
    def block_size(self) -> int:
        """
        Returns the block size of the algorithm in bytes
        """
        if self == EncryptionAlgorithm.Plaintext:
            return 1
        return 16


def encrypt_blob(alg: EncryptionAlgorithm, key: bytes, nonce_iv: bytes, plaintext: Buffer) -> bytes:
    if alg == EncryptionAlgorithm.Plaintext:
        return bytes(plaintext)
    if alg == EncryptionAlgorithm.AES128CTR:
        return encrypt_blob_aes128_ctr(key, nonce_iv, bytes(plaintext))
    elif alg == EncryptionAlgorithm.AES128CBC:
        return encrypt_blob_aes128_cbc(key, nonce_iv, bytes(plaintext))
    elif alg == EncryptionAlgorithm.AES256CTR:
        return encrypt_blob_aes256_ctr(key, nonce_iv, bytes(plaintext))
    elif alg == EncryptionAlgorithm.AES256CBC:
        return encrypt_blob_aes256_cbc(key, nonce_iv, bytes(plaintext))
    raise ValueError(f"Unknown encryption algorithm: {alg}")


def decrypt_blob(alg: EncryptionAlgorithm, key: bytes, nonce_iv: bytes, ciphertext: Buffer) -> bytes:
    if alg == EncryptionAlgorithm.Plaintext:
        return bytes(ciphertext)
    if alg == EncryptionAlgorithm.AES128CTR:
        return decrypt_blob_aes128_ctr(key, nonce_iv, bytes(ciphertext))
    elif alg == EncryptionAlgorithm.AES128CBC:
        return decrypt_blob_aes128_cbc(key, nonce_iv, bytes(ciphertext))
    elif alg == EncryptionAlgorithm.AES256CTR:
        return decrypt_blob_aes256_ctr(key, nonce_iv, bytes(ciphertext))
    elif alg == EncryptionAlgorithm.AES256CBC:
        return decrypt_blob_aes256_cbc(key, nonce_iv, bytes(ciphertext))
    raise ValueError(f"Unknown encryption algorithm: {alg}")


def decrypt_blob_aes128_ctr(key: bytes, nonce: bytes, ciphertext: bytes) -> bytes:
    unpadder = padding.PKCS7(algorithms.AES128.block_size).unpadder()
    cipher = Cipher(
        algorithm=algorithms.AES128(key),
        mode=modes.CTR(nonce),
        backend=backends.default_backend(),
    )
    decryptor = cipher.decryptor()
    plaintext = decryptor.update(ciphertext) + decryptor.finalize()
    return unpadder.update(plaintext) + unpadder.finalize()


def decrypt_blob_aes256_ctr(key: bytes, nonce: bytes, ciphertext: bytes) -> bytes:
    unpadder = padding.PKCS7(algorithms.AES256.block_size).unpadder()
    cipher = Cipher(
        algorithm=algorithms.AES256(key),
        mode=modes.CTR(nonce),
        backend=backends.default_backend(),
    )
    decryptor = cipher.decryptor()
    plaintext = decryptor.update(ciphertext) + decryptor.finalize()
    return unpadder.update(plaintext) + unpadder.finalize()


def decrypt_blob_aes128_cbc(key: bytes, iv: bytes, ciphertext: bytes) -> bytes:
    unpadder = padding.PKCS7(algorithms.AES128.block_size).unpadder()
    cipher = Cipher(
        algorithm=algorithms.AES128(key),
        mode=modes.CBC(iv),
        backend=backends.default_backend(),
    )
    decryptor = cipher.decryptor()
    plaintext = decryptor.update(ciphertext) + decryptor.finalize()
    return unpadder.update(plaintext) + unpadder.finalize()


def decrypt_blob_aes256_cbc(key: bytes, iv: bytes, ciphertext: bytes) -> bytes:
    unpadder = padding.PKCS7(algorithms.AES256.block_size).unpadder()
    cipher = Cipher(
        algorithm=algorithms.AES256(key),
        mode=modes.CBC(iv),
        backend=backends.default_backend(),
    )
    decryptor = cipher.decryptor()
    plaintext = decryptor.update(ciphertext) + decryptor.finalize()
    return unpadder.update(plaintext) + unpadder.finalize()


def encrypt_blob_aes128_ctr(key: bytes, nonce: bytes, plaintext: bytes) -> bytes:
    padder = padding.PKCS7(algorithms.AES128.block_size).padder()
    cipher = Cipher(
        algorithm=algorithms.AES128(key),
        mode=modes.CTR(nonce),
        backend=backends.default_backend(),
    )
    encryptor = cipher.encryptor()
    plaintext = padder.update(plaintext) + padder.finalize()
    return encryptor.update(plaintext) + encryptor.finalize()


def encrypt_blob_aes256_ctr(key: bytes, nonce: bytes, plaintext: bytes) -> bytes:
    padder = padding.PKCS7(algorithms.AES256.block_size).padder()
    cipher = Cipher(
        algorithm=algorithms.AES256(key),
        mode=modes.CTR(nonce),
        backend=backends.default_backend(),
    )
    encryptor = cipher.encryptor()
    plaintext = padder.update(plaintext) + padder.finalize()
    return encryptor.update(plaintext) + encryptor.finalize()


def encrypt_blob_aes128_cbc(key: bytes, iv: bytes, plaintext: bytes) -> bytes:
    padder = padding.PKCS7(algorithms.AES128.block_size).padder()
    cipher = Cipher(
        algorithm=algorithms.AES128(key),
        mode=modes.CBC(iv),
        backend=backends.default_backend(),
    )
    encryptor = cipher.encryptor()
    plaintext = padder.update(plaintext) + padder.finalize()
    return encryptor.update(plaintext) + encryptor.finalize()


def encrypt_blob_aes256_cbc(key: bytes, iv: bytes, plaintext: bytes) -> bytes:
    padder = padding.PKCS7(algorithms.AES256.block_size).padder()
    cipher = Cipher(
        algorithm=algorithms.AES256(key),
        mode=modes.CBC(iv),
        backend=backends.default_backend(),
    )
    encryptor = cipher.encryptor()
    plaintext = padder.update(plaintext) + padder.finalize()
    return encryptor.update(plaintext) + encryptor.finalize()
